import re
import numpy as np
from astro.coord import s2dec
from matplotlib.mlab import rec_drop_fields,rec_append_fields
from astro.io import readtxt

def readheader(filename):
    """ Read in the header of an aaomega text file.

    Returns the header as a dictionary
    """
    fh = open(filename,'r')
    info = []
    info.append(fh.next())
    while info[-1].lstrip().startswith('#'):
        info.append(fh.next())
    info = info[:-1]
    fh.close()
    dinfo = dict()
    for line in info:
        line = line.lstrip('# ')
        key = line.split(None,1)[0]
        if key in ('note', 'spline_points'):
            val = line[line.find(key) + len(key):]
        else:
            val = line.split()[1]
        dinfo[key] = val
    return dinfo


def rec_convert_field(rec, name, func, invalid=None):
    """ Converts each entry in a field in the record array with the
    given function.  Returns the modified array. If invalid is given,
    it uses that value if the function fails when trying to convert an
    entry."""
    # try to convert all the elements
    temp = []
    for item in rec[name]:
        try:
            temp.append(func(item))
        except ValueError:
            if invalid is not None:
                temp.append(invalid)
            else:
                raise

    newrec = rec_drop_fields(rec,[name])
    newrec = rec_append_fields(newrec,name,temp)
    return newrec.view(np.recarray)

def readsex(filename):
    """ Read a sextractor catalogue. """
    fh = open(filename)
    # get the header
    row = fh.next()
    hd = []
    while row.startswith('#'):
        if row[1:].strip():
            hd.append(row)
        row = fh.next()
    fh.close()
    # get column numbers and names
    number, names = zip(*[row.split() for row in hd])[1:3]
    indcol = [int(c)-1 for c in number]
    if len(names) - len(set(names)):
        dup = [n for n in set(names) if names.count(n) > 1]
        raise ValueError('fields with same names: %s' % dup)
    # read in the data
    return readtxt(filename, names=names, usecols=indcol)

def sex2reg(filename, s, colour='green', tag='all'):
    """Create a region file for ds9 to read in from sextractor output file.

    colour is one of: cyan blue magenta red green yellow white black """

    regions = ['global font="helvetica 10 normal" select=1 highlite=1 '
               'edit=0 move=1 delete=1 include=1 fixed=0 source']
    regions.append('image')

    fields = 'X_IMAGE Y_IMAGE A_IMAGE B_IMAGE THETA_IMAGE'.split()
    for i,rec in enumerate(s):
        vals = [rec[f] for f in fields]
        vals.extend([i+1, colour, tag])
        format = 'ellipse(%s %s %s %s %s) # text={%s} color=%s tag={%s}'
        regions.append(format % tuple(vals))

    fh = open(filename,'w')
    fh.write('\n'.join(regions))
    fh.close()


def read_veron(filename,sep='|'):
    """ Reads in a veron output file ('|'-separated values). Returns a
    record array. """

    fh = open(filename)
    f = list(fh)
    fh.close()
    # remove blank lines
    f = [r for r in f if r]
    # remove single char lines
    f = [r for r in f if len(r) > 1]
    # get header
    hd = [r.split()[1] for r in f if r.startswith('#Column')]
    # remove comment lines
    f = [r for r in f if not r.startswith('#')]
    print 'reading in Veron catalogue...'
    qsos = np.rec.fromarrays(readtxt(f[3:],sep=sep,arrays=0),names=','.join(hd))
    qsos = rec_convert_field(qsos,'z',float, invalid=-1)

    ra,dec = zip(*[s2dec(q['RAJ2000'],q['DEJ2000']) for q in qsos])
    
    qsos = rec_append_fields(qsos,'ra',ra)
    qsos = rec_append_fields(qsos,'dec',dec)

    return qsos

def read_simbad(filename):
    """ Reads in a simbad output file (ascii '|'-separated file). Note
    that ra/dec should be in decimal format (You need to change the
    default simbad output options).

    Returns a record array."""

    temp = open(filename)
    f = temp.readlines()
    temp.close()
    # remove blank lines
    f = [r for r in f if r]
    # remove single char lines
    f = [r for r in f if len(r) > 1]
    # get header
    hd, = [r for r in f if r.startswith(' #')]    
    hd = [re.sub('[(),./\\\\]',' ',item).strip() for item in hd.split('|')]
    hd = [item.replace(' ','_').replace('#','no') for item in hd]
    
    # select object lines
    f = [r for r in f if r[0].isdigit()]
    print 'reading in Simbad catalogue...'
    qsos = np.rec.fromarrays(readtxt(f,sep='|',arrays=0),names=','.join(hd))

    ra,dec = zip(*[[float(i) for i in q['coord1__ICRS_2000_2000'].split()]
                 for q in qsos])
    qsos = rec_append_fields(qsos,'ra',ra)
    qsos = rec_append_fields(qsos,'dec',dec)

    qsos = rec_convert_field(qsos,'redshift',float, invalid=-1)
    qsos = rec_append_fields(qsos,'z',qsos['redshift'])
    qsos = rec_drop_fields(qsos,['redshift'])

    # remove absorption line entries
    qsos = qsos[(qsos.typ != 'ALS') & (qsos.typ != 'DLA') & 
                (qsos.typ != 'LyA') & (qsos.typ != 'mAL') &
                (qsos.typ != 'LLS') & (qsos.typ != 'BAL')]

    return qsos

def read_ned(filename):
    """ Reads in a NED output file (ascii '|'-separated file). Note
    that ra/dec should be in decimal format (You need to change the
    default simbad output options).

    Returns a record array."""

    temp = open(filename)
    f = temp.readlines()
    temp.close()
    # remove blank lines
    f = [r for r in f if r]
    # remove single char lines
    f = [r for r in f if len(r) > 1]
    # get header
    hd, = [r for r in f if r.startswith('No.')]    
    hd = [re.sub('[(),./\\\\]',' ',item).strip() for item in hd.split('|')]
    hd = [item.replace(' ','_') for item in hd]
    
    # select object lines
    f = [r for r in f if r[0].isdigit()]
    print 'reading in NED catalogue...'
    qsos = np.rec.fromarrays(readtxt(f,sep='|',arrays=0),names=','.join(hd))

    qsos = rec_convert_field(qsos,'Redshift',float, invalid=-1)
    qsos = rec_append_fields(qsos,'z',qsos['Redshift'])
    qsos = rec_append_fields(qsos,'ra',qsos['RA_deg'].astype(float))
    qsos = rec_append_fields(qsos,'dec',qsos['DEC_deg'].astype(float))
    qsos = rec_drop_fields(qsos,['RA_deg','DEC_deg','Redshift'])

    # remove absorption lines
    qsos = qsos[(qsos.Type != 'AbLS')]

    return qsos



def read_ned_batch(filename):
    """ Reads in a NED batch request result file (see
    http://nedwww.ipac.caltech.edu/help/batch.html), extracting
    magnitude, redshift, ra, dec, name, etc.  Returns record array.
    """
    temp = open(filename)
    f = temp.read()
    temp.close()

    # find start and end of results:
    i1 = f.index('************************** SEARCH RESULTS '
                 '******************************\n')
    i2 = f.index('    **********  Names for all the objects found '
                 'in this request **********\n')
    f = f[i1:i2]
    # separate results
    results = f.split('++++++++++++++++++++++++++++ NEXT SEARCH '
                      '+++++++++++++++++++++++++++++++++\n')
    results = [r.split('\n') for r in results]

    info = []
    # for each result, record redshift, mag, name, ra/dec
    for r in results:
        # find the name, type, position:
        try:
            i1 = r.index('#--Object Name------------Type-------------Position'
                         '---------Dist.-Ref-Note-Phot')
        except ValueError:
            # if it wasn't found:
            info.append([-1]*10)
            continue
        row = r[i1+1]
        name = row[1:26].strip()
        nedtype = row[26:33].strip()
        sra = row[33:47].strip()
        sdec = row[47:60].strip()
        ra,dec = s2dec(sra,sdec)
        # now redshift, mag, class
        i1 = r.index('#           Extinct    --------column1------- column2'
                     ' column3 column4 -----------------column5'
                     '---------------- ')
        row = r[i1+1]
        mag = row[46:54].strip()
        if mag == '':
            mag = -1.0
            magtype = '-1'
        else:
            if mag[-1].isalpha():
                magtype = mag[-1]
                mag = float(mag[:-1])
            else:
                magtype = '-1'
                mag = float(mag)
        nedclass = ','.join(row[23:46].split())
        #print row[70:]
        redshift = row[70:].strip()
        if redshift == '':
            redshift = -1
        else:
            try: 
                redshift = float(redshift.split()[0])
            except ValueError:
                redshift = -1
        info.append((name,nedtype,nedclass,ra,dec,sra,sdec,mag,magtype,
                     redshift))

    return np.rec.fromrecords(info,names='name,nedtype,nedclass,ra,dec,sra'
                              ',sdec,mag,magtype,z')

